# encoding: utf-8

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import math

class GrayForecast():
    def __init__(self, X: pd.core.frame.DataFrame, dataCol=None):
        if isinstance(X, pd.core.frame.DataFrame):
            self.X = X
            self.X.columns = ['数据']
        elif isinstance(X, pd.core.series.Series):
            self.X = pd.DataFrame(X)
            self.X.columns = ['数据']
        else:
            self.X = pd.DataFrame(X)
            self.X.columns = ['数据']
        self.Xc = self.X.copy()
        if dataCol:
            self.datacolumn = dataCol
        else:
            self.datacolumn = None

    # 数据级比校验
    def level_check(self):
        n = len(self.X)
        lambda_k = np.zeros(n-1)
        for i in range(n-1):
            lambda_k[i] = self.X.ix[i]/self.X.ix[i+1]
            if lambda_k[i] < np.exp(-2/(n+1)) or lambda_k[i] > np.exp(2/(n+2)):
                flag = False
        else:
            flag = True

        self.lambda_k = lambda_k

        if not flag:
            print("级比校验失败，请对X(0)做平移变换")
            return False
        else:
            # print("级比校验成功！")
            return True
            
    # GM(1,1)
    def GM_11_build_model(self, forecast=5):
        if forecast > len(self.X):
            raise Exception('您的数据行不够')
        X_0 = np.array(self.Xc.tail(forecast))
        # 1-AGO
        X_1 = np.zeros(X_0.shape)
        for i in range(X_0.shape[0]):
            X_1[i] = np.sum(X_0[0:i+1])
        # 紧邻均值生成序列
        Z_1 = np.zeros(X_1.shape[0]-1)
        for i in range(1, X_1.shape[0]):
            Z_1[i-1] = -0.5*(X_1[i]+X_1[i-1])

        B = np.append(np.array(np.mat(Z_1).T), np.ones(Z_1.shape).reshape((Z_1.shape[0], 1)), axis=1)
        Yn = X_0[1:].reshape((X_0[1:].shape[0], 1))

        B = np.mat(B)
        Yn = np.mat(Yn)
        a_ = np.linalg.inv(B.T.dot(B)).dot(B.T).dot(Yn)

        a, b = np.array(a_.T)[0]
        # print(a,b);input()

        X_ = np.zeros(X_0.shape[0])
        def f(k):
            return (X_0[0]-b/a)*(1-np.exp(a))*np.exp(-a*(k))

        self.Xc.loc[len(self.Xc)] = f(X_.shape[0])

    # 预测模型
    def forecast(self, time=5, forecast_data_len=None):
        if forecast_data_len == None:
            forecast_data_len = len(self.X.to_numpy().flatten())
        if self.level_check():
            for i in range(time):
                self.GM_11_build_model(forecast=forecast_data_len)
            return self.Xc

    # 打印当前预测序列
    def log(self):
        res = self.Xc.copy()
        if self.datacolumn:
            res.columns = [self.datacolumn]
        return res
    
    # reset初始化序列
    def reset(self):
        self.Xc = self.X.copy()

    # 作图
    def plot(self):
        self.Xc.plot()
        if self.datacolumn:
            plt.ylabel(self.datacolumn)
            plt.legend([self.datacolumn])

    
def grayforecast2(X):
    history_data = X.values()
    n = len(history_data)
    X0 = np.array(history_data)

    #累加生成
    history_data_agg = [sum(history_data[0:i+1]) for i in range(n)]
    X1 = np.array(history_data_agg)

    #计算数据矩阵B和数据向量Y
    B = np.zeros([n-1,2])
    Y = np.zeros([n-1,1])
    for i in range(0,n-1):
        B[i][0] = -0.5*(X1[i] + X1[i+1])
        B[i][1] = 1
        Y[i][0] = X0[i+1]

    #计算GM(1,1)微分方程的参数a和u
    #A = np.zeros([2,1])
    A = np.linalg.inv(B.T.dot(B)).dot(B.T).dot(Y)
    a = A[0][0]
    u = A[1][0]

    #建立灰色预测模型
    XX0 = np.zeros(n)
    XX0[0] = X0[0]
    for i in range(1,n):
        XX0[i] = (X0[0] - u/a)*(1-math.exp(a))*math.exp(-a*(i));

    #模型精度的后验差检验
    e = 0      #求残差平均值
    for i in range(0,n):
        e += (X0[i] - XX0[i])
    e /= n

    #求历史数据平均值
    aver = 0;     
    for i in range(0,n):
        aver += X0[i]
    aver /= n

    #求历史数据方差
    s12 = 0;     
    for i in range(0,n):
        s12 += (X0[i]-aver)**2;
    s12 /= n

    #求残差方差
    s22 = 0;       
    for i in range(0,n):
        s22 += ((X0[i] - XX0[i]) - e)**2;
    s22 /= n

    #求后验差比值
    C = s22 / s12   

    #求小误差概率
    cout = 0
    for i in range(0,n):
        if abs((X0[i] - XX0[i]) - e) < 0.6754*math.sqrt(s12):
            cout = cout+1
        else:
            cout = cout
    P = cout / n

    if (C < 0.35 and P > 0.95):
        #预测精度为一级
        m = 3   #请输入需要预测的年数
        #print('往后m各年负荷为：')
        f = np.zeros(m)
        for i in range(0,m):
            f[i] = (X0[0] - u/a)*(1-math.exp(a))*math.exp(-a*(i+n))    
        print(f)
    else:
        print('灰色预测法不适用')